from datetime import timedelta
from pathlib import Path
from timeit import timeit

import numpy as np

from raysect.core.math.function.float.function3d.interpolate import Discrete3DMesh
from raysect.primitive.mesh.tetra_mesh import TetraMeshData


def load_mesh(mesh_filepath):
    vertices = []
    tetrahedra = []

    with open(mesh_filepath, "r") as f:
        lines = f.readlines()

    i = 0
    while i < len(lines):
        line = lines[i].strip()
        if not line or line.startswith("#"):
            i += 1
            continue

        # Check section headers (case-sensitive)
        if line == "Vertices":
            i += 1
            vertex_count = int(lines[i].strip())
            i += 1
            for _ in range(vertex_count):
                parts = lines[i].strip().split()
                # take only the first 3 coordinates
                vertices.append([float(v) for v in parts[:3]])
                i += 1
            continue
        elif line == "Tetrahedra":
            i += 1
            tet_count = int(lines[i].strip())
            i += 1
            for _ in range(tet_count):
                parts = lines[i].strip().split()
                # take only the first 4 indices, converting them to int.
                tetrahedra.append([int(idx) for idx in parts[:4]])
                i += 1
            continue

        i += 1

    vertices = np.array(vertices)
    tetrahedra = np.array(tetrahedra, dtype=np.int32) - 1  # 0-based indexing
    return vertices, tetrahedra


def create_discrete3dmesh(vertices, tetrahedra):
    return Discrete3DMesh(
        vertices, tetrahedra, np.ones((tetrahedra.shape[0])), False, 0
    )


def create_tetrameshdata(vertices, tetrahedra):
    return TetraMeshData(vertices, tetrahedra)


def load_tetrameshdata(mesh_filepath):
    return TetraMeshData.from_file(mesh_filepath)


if __name__ == "__main__":
    ROOT = Path(__file__).parent
    mesh_file = ROOT / "demos" / "resources" / "stanford_bunny.mesh"
    vertices, tetrahedra = load_mesh(mesh_file)

    # Save tetramesh data in advance
    tetra = TetraMeshData(vertices, tetrahedra)
    tetra_file = ROOT / "temp_tetra.rsm"
    tetra.save(tetra_file)

    # === Measure time ===
    loop = 5

    # create Discrete3DMesh
    result = timeit(
        "create_discrete3dmesh(vertices, tetrahedra)", globals=globals(), number=loop
    )
    elapsed_time = timedelta(seconds=result / loop)
    print(f"Create Discrete3DMesh: {elapsed_time}")

    # create TetraMeshData
    result = timeit(
        "create_tetrameshdata(vertices, tetrahedra)", globals=globals(), number=loop
    )
    elapsed_time = timedelta(seconds=result / loop)
    print(f"Create TetraMeshData: {elapsed_time}")

    # load TetraMeshData from file
    result = timeit("load_tetrameshdata(tetra_file)", globals=globals(), number=loop)
    elapsed_time = timedelta(seconds=result / loop)
    print(f"Load TetraMeshData from file: {elapsed_time}")